package io.gatling.extensions.kafka.action

import com.typesafe.scalalogging.StrictLogging
import io.gatling.commons.stats.{KO, OK}
import io.gatling.commons.util.Clock
import io.gatling.core.CoreComponents
import io.gatling.core.action.{Action, ExitableAction}
import io.gatling.core.session.{Expression, Session}
import io.gatling.core.stats.StatsEngine
import io.gatling.core.util.NameGen
import io.gatling.extensions.kafka.protocol.KafkaProducerComponents
import org.apache.kafka.clients.producer.{ProducerRecord, RecordMetadata}

case class KafkaProductAction(
                               name: String,
                               topic: Expression[String],
                               key: Option[Expression[String]],
                               value: Expression[String],
                               kafkaProducerComponents: KafkaProducerComponents,
                               coreComponents: CoreComponents,
                               throttled: Boolean,
                               next: Action,
                             ) extends ExitableAction with NameGen with StrictLogging {

  import coreComponents._

  override def statsEngine: StatsEngine = coreComponents.statsEngine

  override def clock: Clock = coreComponents.clock

  override def execute(session: Session): Unit = {

    val start: Long = clock.nowMillis

    key match {
      case Some(keyExpression) =>
        kafkaProducerComponents.producerMap.get(session.userId).send(new ProducerRecord(topic(session).toOption.get, keyExpression(session).toOption.get, value(session).toOption.get), (_: RecordMetadata, e: Exception) => if (e == null) success(session, start) else fail(session, e.getMessage, start))
      case None =>
        // 这里 send 方法没有错误，编译器识别错了
        kafkaProducerComponents.producerMap.get(session.userId).send(new ProducerRecord(topic(session).toOption.get, value(session).toOption.get), (_: RecordMetadata, e: Exception) => if (e == null) success(session, start) else fail(session, e.getMessage, start))
    }

    throttler match {
      case Some(th) if throttled => th.throttle(session.scenario, () => next ! session)
      case None => next ! session
    }
  }

  def fail(session: Session, message: String, start: Long): Unit = {
    statsEngine.logResponse(session.scenario, session.groups, name, start, clock.nowMillis, KO, Some(message), Some(message))
  }

  def success(session: Session, start: Long): Unit = {
    statsEngine.logResponse(session.scenario, session.groups, name, start, clock.nowMillis, OK, None, None)
  }
}
